from typing import Literal

from qtpy.QtSql import QSqlDatabase, QSqlQuery, QSql
from pymysql import Connection


# connect = Connection()
# connect.connect()
# cursor = connect.cursor()
# cursor.execute()


class DatabaseManager:
    def __init__(
            self,
            driver: Literal["QSQLITE", "QMYSQL", "QPSQL", "QODBC", "QODBCDriver", ...],
            connectionName: str,
            database: str,
            user: str = '',
            password: str = '',
            host: str = 'localhost',
            port: int = 3306,
    ):
        self.db = QSqlDatabase.addDatabase(driver, connectionName)
        self.query = QSqlQuery(self.db)
        self.db.setDatabaseName(database)

        if driver == "QSQLITE" or driver == "QODBC" or driver == "QODBCDriver":
            pass
        else:
            self.db.setHostName(host)
            self.db.setPort(port)
            self.db.setUserName(user)
            self.db.setPassword(password)

    def __del__(self):
        self.close()

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def _get_all_table_data(self, query: QSqlQuery):
        data = []
        while query.next():
            row = {}
            for i in range(query.record().count()):
                field_name = query.record().fieldName(i)
                field_value = query.value(i)
                row[field_name] = field_value
            data.append(row)
        return data

    def table_info(self, table_name: str):
        if self.driver() == "QSQLITE":
            query = self.db.exec(f"PRAGMA table_info({table_name});")
        else:
            query = self.db.exec(f"DESCRIBE {table_name};")
        return self._get_all_table_data(query)

    def table_columns(self, table_name: str):
        query = self.db.exec(f"SELECT * FROM {table_name} LIMIT 1;")
        return [query.record().fieldName(i) for i in range(query.record().count())]

    def tables(self, type: QSql.TableType = QSql.TableType.AllTables):
        return self.db.tables(type)

    def connect(self):
        success = self.db.open()
        if not success:
            raise ConnectionError("Could not connect to database")
        return success

    def select(self, query: str) -> list:
        return self._get_all_table_data(self.db.exec(query))

    def insert(self, query: str) -> bool:
        print(query)
        self.query.prepare(query)
        return self.query.exec()

    def rollback(self):
        self.db.rollback()

    def commit(self):
        self.db.commit()

    def close(self):
        self.query.finish()
        self.query.clear()
        self.db.close()

    def isOpen(self):
        return self.db.isOpen()

    def driver(self):
        return self.db.driverName()


class SQLiteDatabaseManager(DatabaseManager):
    def __init__(
            self,
            connectionName: str = "default",
            database: str = "./test.db",
    ):
        super().__init__(
            driver="QSQLITE", connectionName=connectionName, database=database)


if __name__ == '__main__':
    import sys
    from qtpy.QtWidgets import QApplication

    app = QApplication(sys.argv)
    # manager = SQLiteDatabaseManager()
    # manager.connect()
    # table = manager.execute('SELECT * FROM users')
    # while table.next():
    #     print(table.value(0))
    # manager.close()
    with SQLiteDatabaseManager() as manager:
        users = manager.insert(f'INSERT INTO users (name, phone) VALUES {"John","13512345678"}')
        # users = manager.insert('SELECT * FROM users')
        print(users)
    sys.exit(app.exec_())
